import * as tf from '@tensorflow/tfjs'

function createSimpleModel(index, activation){
  var model = ''
  console.log(index,activation)
  if(index === 3){
    // 创建3层全连接网络
    model = tf.sequential();
    model.add(tf.layers.dense({units: 16, inputShape:[28, 28, 3]}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 6, activation: 'softmax'}));
    // console.log(model);

  }else if(index === 4){
    //创建4层全连接网络
    model = tf.sequential();
    model.add(tf.layers.dense({units: 16, inputShape:[28, 28, 3]}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 6, activation: 'softmax'}));
    console.log(model);
  }else{
    //创建5层全连接网络
    model = tf.sequential();
    model.add(tf.layers.dense({units: 16, inputShape:[28, 28, 3]}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.dense({units: 32, activation:  activation}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 6, activation: 'softmax'}));
    console.log(model);
    // console.log(model.summary());
  }
  return model;
}

function createComplexModel(index, activation){
  var model = ''
  if(index === 3){
    //创建3层卷积神经网络
    model = tf.sequential();
    model.add(tf.layers.conv2d({inputShape: [28, 28, 3], kernelSize: 3, strides: 1, filters: 1, activation: activation, padding:'same'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 5, strides: 1, padding:'same', filters: 3, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 5, strides: 1, padding:'same', filters: 3, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 6, activation: 'softmax'}));
  }else if(index === 4){
    //创建4层卷积神经网络
    model = tf.sequential();
    model.add(tf.layers.conv2d({inputShape: [28, 28, 3], kernelSize: 3, strides: 1, filters: 1, padding:'same', activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 3, filters: 3, padding:'same', strides: 1, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 5, filters: 6, padding:'same', strides: 1, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 7, filters: 12, padding:'same', strides: 1, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 6, activation: 'softmax'}));
  }else{
    //创建5层卷积神经网络
    model = tf.sequential();
    model.add(tf.layers.conv2d({inputShape: [28, 28, 3], kernelSize: 3, strides: 1, filters: 1,  padding:'same', activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 3, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 6, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 12, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 24, activation: activation}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 6, activation: 'softmax'}));
  }
  return model;
}

export{
  createSimpleModel,
  createComplexModel
}
